import os

import utils
from config import ENVIRONMENT_CONFIG


class Account:
    '''
    账户类：存有账户所需的必要信息
    '''
    def __init__(self,
                 cash=1000000,
                 skid_rate=0,
                 fee_rate=0):
        self.cash = cash
        self.equity = self.cash
        self.skid_rate = skid_rate
        self.fee_rate = fee_rate
        self.position = {} # 仓位
        self.order_id = 0#订单编号

        utils.csv_write(
            os.path.join(ENVIRONMENT_CONFIG['output_path'], 'orders.csv'),
            ['order_id', 'trade_date', 'code', 'buyorsell', 'price', 'vol', 'amt'],
            new=True
        )

        utils.csv_write(
            os.path.join(ENVIRONMENT_CONFIG['output_path'], 'wallet.csv'),
            ['trade_date', 'cash', 'stocks_value', 'equity', 'position'],
            new=True
        )


    def buy_order(self, code, price, vol, d):
        '''
        买入若干数量这一支股票
        :param code: 股票代号
        :param price: 单股价格
        :param vol: 股票数量
        :param d: 交易日期
        :return:
        '''

        price = round(price * (1+self.skid_rate), 2)
        amt = int(price * vol)#成交金额
        fee = round(amt * self.fee_rate, 2)

        if self.cash >= amt + fee:
            self.cash -= amt +fee
            self.position.update({code: int(self.position.get(code, 0) + vol)})

            utils.csv_write(
                os.path.join(ENVIRONMENT_CONFIG['output_path'], 'orders.csv'),
                [str(self.order_id).rjust(8, '0'), d, code, 'buy', str(price), str(vol), str(amt)]
            )
            self.order_id += 1
        else:
            pass


    def sell_order(self, code, price, vol, d):
        '''
        卖出若干数量这支股票
        :param code: 股票代号
        :param price: 单股价格
        :param vol: 股票数量
        :param d: 交易日期
        :return:
        '''
        price *= round(1-self.skid_rate, 2)
        amt = int(price * vol)
        fee = round(amt * self.fee_rate, 2)

        self.cash += amt - fee
        self.position.update({code: self.position.get(code, 0) - vol})

        utils.csv_write(
            os.path.join(ENVIRONMENT_CONFIG['output_path'], 'orders.csv'),
            [str(self.order_id).rjust(8, '0'), d, code, 'sell', str(price), str(vol), str(amt)]
        )
        self.order_id += 1

    def update_equity(self, prices_data):#更新资产
        '''
        跟新资产，因为股票价格每日都会变动
        :param prices_data: 开盘价字典
        :return:
        '''
        stocks_value = 0
        for stock in self.position:
            stocks_value += round(self.position.get(stock, 0) * prices_data.get(stock, 0), 2)
        self.equity = round(self.cash + stocks_value, 2)

    def update_wallet(self, d):#更新账户信息
        '''
        把卖出的股票清除
        :param d: 日期
        :return:
        '''
        for code in list(self.position.keys()):
            if self.position[code] == 0:
                del self.position[code]
        utils.csv_write(
            os.path.join(ENVIRONMENT_CONFIG['output_path'], 'wallet.csv'),
            [d, str(round(self.cash, 2)), str(int(self.equity-self.cash)), str(self.equity), str(self.position)],
        )
